package com.shujia.flink.state

import java.util.Properties

import org.apache.flink.api.common.functions.{ReduceFunction, RichMapFunction, RuntimeContext}
import org.apache.flink.api.common.serialization.SimpleStringSchema
import org.apache.flink.api.common.state._
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer

object Demo2ReducingState {
  def main(args: Array[String]): Unit = {

    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment


    val properties = new Properties()
    properties.setProperty("bootstrap.servers", "master:9092,node1:9092,node2:9092")
    //properties.setProperty("group.id", "asdasdsad")
    //properties.setProperty("auto.offset.reset", "earliest") //earliest 读取所有数据、latest ： 度最新数据


    val consumer = new FlinkKafkaConsumer[String]("test_topic1", new SimpleStringSchema(), properties)

    consumer.setStartFromLatest()


    val kafkaDS: DataStream[String] = env.addSource(consumer)


    val kvDS: KeyedStream[(String, Int), String] = kafkaDS.map((_, 1)).keyBy(_._1)


    //通过聚合状态统计单词的数量
    val countDS: DataStream[(String, Int)] = kvDS.map(new MyRichMapFunction)


    countDS.print()


    env.execute()

  }
}


class MyRichMapFunction extends RichMapFunction[(String, Int), (String, Int)] {


  var reduceState: ReducingState[Int] = _


  override def open(parameters: Configuration): Unit = {
    val context: RuntimeContext = getRuntimeContext


    /**
      * 聚合状态的描述对象，需要一个聚合函数
      *
      */
    val reducingStateDescriptor = new ReducingStateDescriptor[Int]("reduce", new ReduceFunction[Int] {
      override def reduce(x: Int, y: Int): Int = {
        x + y
      }
    }, classOf[Int])


    reduceState = context.getReducingState(reducingStateDescriptor)


  }

  override def map(value: (String, Int)): (String, Int) = {
    //累加计算
    //更新单词的数量
    reduceState.add(value._2)

    //获取聚合结果
    //获取最新单词的数量
    val count: Int = reduceState.get()

    //返回结果
    (value._1, count)
  }
}
